# -*- coding: utf-8 -*-
"""
Created on Thu Apr 21 23:33:12 2022

@author: 13503
"""

import torch

from AETrainer import AETrainer
from SignNetData import SignNetData
from StackedAutoEncoder import StackedAutoEncoder

class SAETrainer():
    def __init__(self, sae:StackedAutoEncoder, data_path:str):
        self.models = sae
        self.data_path = data_path

    def train(self, epochs:int, mode:str="train"):
        sign_net = SignNetData(file_path=self.data_path, train_ratio=0.8)
        if mode == "train":
            rep_path = sign_net.train_mat_save_path
        elif mode == "validate":
            rep_path = sign_net.validate_mat_save_path
        elif mode == "test":
            rep_path = sign_net.test_mat_save_path
        else:
            raise "Error! Argument mode is illegal!"
        
        self.data_path = rep_path
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        penalty = torch.load(sign_net.penalty_mat_save_path).to(device)

        for model in self.models.sae:
            print("<<<<<<<<<< begin to train model: {0} >>>>>>>>>>".format(model.model_name))
            ae_trainer =  AETrainer(model, rep_path, penalty)
            ae_trainer.train(epochs)
            rep_path = ae_trainer.save_rep_path
        return rep_path